import gradio as gr
import numpy as np
from PIL import Image
from pinecone import Pinecone

# 初始化 Pinecone
pinecone = Pinecone(api_key="543dff37-939d-49f8-997d-45a361e9abfb")
index_name = "mnist-index"

# 检查索引是否存在
if index_name not in pinecone.list_indexes().names():
    raise ValueError(f"索引 '{index_name}' 不存在，请检查索引名称。")

index = pinecone.Index(index_name)

# 定义预测函数
def predict_digit(image_dict):
    # 从 image_dict 中获取图像 composite 部分
    composite_image = image_dict['composite']
    
    # 将图像从字典形式转化为PIL格式
    image = Image.fromarray(composite_image)
    
    # 提取alpha通道，并将其转换为灰度图像
    alpha_channel = image.getchannel('A')
    
    # 调整图像大小为8x8，因为加载的模型是基于这个大小训练的
    resized_image = alpha_channel.resize((8, 8), Image.Resampling.LANCZOS) 
    
    # 转换为NumPy数组，并缩放像素值到与训练数据匹配的范围（0-16）
    image_array = np.array(resized_image)
    scaled_image = (image_array / 255.0) * 16
    
    # 展平图像为1D数组，符合模型的输入要求
    flattened_image = scaled_image.flatten().tolist()
    
    # 使用 Pinecone 的查询功能进行推理
    query_result = index.query(vector=flattened_image, top_k=1, include_metadata=True)
    
    # 返回预测结果
    prediction = query_result['matches'][0]['metadata']['label']
    return int(prediction)

# 创建Gradio接口
canvas = gr.Sketchpad(canvas_size=(256, 256), label="绘制数字")
interface = gr.Interface(
    fn=predict_digit,
    inputs=canvas,
    outputs="label",
    title="手写数字识别",
    description="请在画布上绘制一个0-9之间的数字，并点击提交进行预测。"
)

# 启动Gradio接口，用户可以通过这个接口进行交互
if __name__ == "__main__":
    interface.launch(share=True)